from netualbuild import NeuralNetwork
from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from torch import nn
import torch
import visdom
import gradio as gr

vis = visdom.Visdom(env="fashion_env")
learning_rate = 1e-3
batch_size = 64
loss_fn = nn.CrossEntropyLoss()

device=(
      "cuda:0"
      if torch.cuda.is_available()
      else "mps"
      if torch.backends.mps.is_available()
      else "cpu"
)

print(f"device: {device}")
model = NeuralNetwork().to(device)
print(model)
optmizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

def update_model(learning_rate, batch_size):
    for param_group in optmizer.param_groups():
        param_group["lr"] = learning_rate

iface = gr.Interface(
    fn=update_model,
    inputs=[gr.Slider(0.001,0.1), gr.Number(32)],
    outputs="text"
)
iface.launch()

def main():
  labels_map = {
      0: "T-Shirt",
      1: "Trouser",
      2: "Pullover",
      3: "Dress",
      4: "Coat",
      5: "Sandal",
      6: "Shirt",
      7: "Sneaker",
      8: "Bag",
      9: "Ankle Boot",
  }

  training_data = datasets.FashionMNIST(
     root="data", train=True, download=True, transform=ToTensor())
  test_data = datasets.FashionMNIST(
     root="data", train=False, download=True, transform=ToTensor())
  train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
  test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

  # optmizer.param_groups()
  # optmizer = torch.optim.SGD(model.parameters(),lr=learning_rate)

  for epochindex in range(0, 50):
    model.train()
    for batchindex, (image, label) in enumerate(train_dataloader):
      image = image.to(device)
      label = label.to(device)
      pred = model(image)
      loss = loss_fn(pred, label)
      # if batchindex % 100 == 0:
      # print(f"batchindex:{batchindex} loss:{loss.item()}\n")
      loss.backward()
      optmizer.step()
      optmizer.zero_grad()
    model.eval()
    batchnum = len(test_dataloader)
    datasize = len(test_data)
    test_loss = 0
    test_correct = 0
    with torch.no_grad():
      for batchindex, (image, label) in enumerate(test_dataloader):
        image = image.to(device)
        label = label.to(device)
        pred = model(image)
        test_loss += loss_fn(pred, label)
      # print(f"pred.shape:{pred.shape}  label.shape:{label.shape}\n")
        test_correct += (pred.argmax(1) == label).type(torch.float).sum().item()
      # print(f"pred:{pred.shape} label:{label.shape}, pred[0].size():{pred.size(0)}")
        if torch.equal(pred.argmax(1), label):
          for index in range(0, pred.size(0)):
            predvalue = pred.argmax(1)
          # print(predvalue,label)
            print(f"pred:{labels_map[predvalue[index].item()]} label:{labels_map[label[index].item()]}")
      test_loss /= batchnum
      test_correct /= datasize
      print(f"accuacy:{(100*test_correct):>0.1f}% avgloss:{test_loss:>8f}\n")